import dis
import new

class MissingLabel(Exception):
    """'goto' called without matching 'label'."""
    pass

def goto(fn):
    """
    A function decorator to add the goto command for a function.
    
    Example usage:
    
        >>> @goto
        ... def myfunction(*args, **kwargs):
        ...     counter = 0
        ...     print "A"
        ...
        ...     label .start
        ...
        ...     print "B"
        ...     print "C"
        ...
        ...     if counter == 0:
        ,,,         counter += 1
        ...         goto .start
        ...
        ...     print "D"
        ...     return
        ...
        >>> myfunction()
        A
        B
        C
        B
        C
        D
    
    See <http://code.activestate.com/recipes/576944-the-goto-decorator/>.
    """
    labels = {}
    gotos = {}
    globalName = None
    index = 0
    end = len(fn.func_code.co_code)
    i = 0

    # scan through the byte codes to find the labels and gotos
    while i < end:
        op = ord(fn.func_code.co_code[i])
        i += 1
        name = dis.opname[op]

        if op > dis.HAVE_ARGUMENT:
            b1 = ord(fn.func_code.co_code[i])
            b2 = ord(fn.func_code.co_code[i+1])
            num = b2 * 256 + b1
            
            if name == 'LOAD_GLOBAL':
                globalName = fn.func_code.co_names[num]
                index = i - 1
                i += 2
                continue
                
            if name == 'LOAD_ATTR':
                if globalName == 'label':
                    labels[fn.func_code.co_names[num]] = index
                elif globalName == 'goto':
                    gotos[fn.func_code.co_names[num]] = index
                    
            name = None
            i += 2

    # no-op the labels
    ilist = list(fn.func_code.co_code)
    for label,index in labels.items():
        ilist[index:index+7] = [chr(dis.opmap['NOP'])]*7

    # change gotos to jumps
    for label,index in gotos.items():
        if label not in labels:
            raise MissingLabelError("Missing label: %s"%label)
        
        target = labels[label] + 7   # skip NOPs
        ilist[index] = chr(dis.opmap['JUMP_ABSOLUTE'])
        ilist[index + 1] = chr(target & 255)
        ilist[index + 2] = chr(target >> 8)

    # create new function from existing function
    c = fn.func_code
    newcode = new.code(
        c.co_argcount,
        c.co_nlocals,
        c.co_stacksize,
        c.co_flags,
        ''.join(ilist),
        c.co_consts,
        c.co_names,
        c.co_varnames,
        c.co_filename,
        c.co_name,
        c.co_firstlineno,
        c.co_lnotab)
    newfn = new.function(newcode,fn.func_globals)
    return newfn
    
@goto
def myfunction(*args, **kwargs):
    counter = 0
    print "A"
    label .start
    print "B"
    print "C"
    if counter == 0:
        counter += 1
        goto .start
    print "D"
    return
    
        
